import torch
import torch.nn as nn
import numpy as np
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' )
tgt_vocab_size = 9 # 所有目标词的长度，在10.10.1数据准备这节中计算获得
class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_emb = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        '''
        dec_inputs: [batch_size, tgt_len]
        enc_intpus: [batch_size, src_len]
        enc_outputs: [batsh_size, src_len, d_model] 经过六次encoder之后得到的东西
        '''
        dec_outputs = self.tgt_emb(dec_inputs)
        # [batch_size, tgt_len, d_model]
        # 同样地，对decoder_layer进行词向量的生成
        dec_outputs = \
            self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1).to(device)
        # 计算他的位置向量
        # [batch_size, tgt_len, d_model]

        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).to(device)
        # [batch_size, tgt_len, tgt_len]

        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).to(device)
        # [batch_size, tgt_len, tgt_len]
        # 当前时刻看不到未来时刻的信息

        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + 
                                       dec_self_attn_subsequence_mask), 0).to(device)
        # [batch_size, tgt_len, tgt_len]
        # 布尔+int gt 大于0 为True
        # 把dec_self_attn_pad_mask和dec_self_attn_subsequence_mask里面为True的部分都剔除掉了           #即屏蔽掉了pad也屏蔽掉了mask

        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)  
        # [batc_size, tgt_len, src_len]
        # 在decoder的第二个attention里面使用
        dec_self_attns, dec_enc_attns = [], []

        for layer in self.layers:
            dec_outputs= layer(dec_outputs, 
                               enc_outputs, 
                               dec_self_attn_mask, 
                               dec_enc_attn_mask)
            # dec_outputs: [batch_size, tgt_len, d_model],
        return dec_outputs